"""
Kalman滤波器实现模块

该模块实现了用于目标跟踪的Kalman滤波器。
主要功能：
1. 状态预测（位置和速度）
2. 观测更新
3. 噪声协方差矩阵管理
4. 状态估计

作者：目标跟踪系统
日期：2024
"""

import numpy as np
from typing import Tuple, Optional


class KalmanFilter:
    """
    Kalman滤波器类
    
    基于线性动态系统的状态估计器，用于预测和更新目标状态。
    状态向量包含位置(x, y)和速度(vx, vy)信息。
    
    参数：
        dt: 时间步长，默认为1.0
        process_noise: 过程噪声标准差，默认为1e-2
        measurement_noise: 测量噪声标准差，默认为1e-1
    """
    
    def __init__(self, dt: float = 1.0, process_noise: float = 1e-2, measurement_noise: float = 1e-1):
        self.dt = dt
        self.process_noise = process_noise
        self.measurement_noise = measurement_noise
        
        # 状态向量维度：[x, y, vx, vy]
        self.state_dim = 4
        self.measurement_dim = 2
        
        # 状态向量和协方差矩阵
        self.state = np.zeros(self.state_dim)  # [x, y, vx, vy]
        self.covariance = np.eye(self.state_dim) * 1000  # 初始协方差矩阵
        
        # 系统矩阵
        self._initialize_matrices()
        
        # 滤波器状态
        self.is_initialized = False
        
    def _initialize_matrices(self) -> None:
        """
        初始化Kalman滤波器的系统矩阵
        
        包括状态转移矩阵、观测矩阵、过程噪声协方差矩阵和测量噪声协方差矩阵。
        """
        # 状态转移矩阵 F (匀速运动模型)
        self.F = np.array([
            [1, 0, self.dt, 0],
            [0, 1, 0, self.dt],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ], dtype=np.float32)
        
        # 观测矩阵 H (只观测位置)
        self.H = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0]
        ], dtype=np.float32)
        
        # 过程噪声协方差矩阵 Q
        self.Q = np.array([
            [self.dt**4/4, 0, self.dt**3/2, 0],
            [0, self.dt**4/4, 0, self.dt**3/2],
            [self.dt**3/2, 0, self.dt**2, 0],
            [0, self.dt**3/2, 0, self.dt**2]
        ], dtype=np.float32) * self.process_noise**2
        
        # 测量噪声协方差矩阵 R
        self.R = np.eye(self.measurement_dim, dtype=np.float32) * self.measurement_noise**2
        
    def initialize(self, initial_position: Tuple[float, float], initial_velocity: Optional[Tuple[float, float]] = None) -> None:
        """
        初始化滤波器状态
        
        设置初始位置和速度，重置协方差矩阵。
        
        参数：
            initial_position: 初始位置 (x, y)
            initial_velocity: 初始速度 (vx, vy)，如果为None则设为0
        """
        if initial_velocity is None:
            initial_velocity = (0.0, 0.0)
            
        self.state = np.array([
            initial_position[0],
            initial_position[1],
            initial_velocity[0],
            initial_velocity[1]
        ], dtype=np.float32)
        
        # 重置协方差矩阵
        self.covariance = np.eye(self.state_dim, dtype=np.float32) * 1000
        
        self.is_initialized = True
        
    def predict(self) -> Tuple[float, float]:
        """
        预测下一时刻的状态
        
        使用状态转移方程预测目标的下一个位置和速度。
        
        返回：
            预测的位置 (x, y)
        """
        if not self.is_initialized:
            raise ValueError("滤波器未初始化，请先调用initialize方法")
            
        # 状态预测
        self.state = self.F @ self.state
        
        # 协方差预测
        self.covariance = self.F @ self.covariance @ self.F.T + self.Q
        
        return (self.state[0], self.state[1])
        
    def update(self, measurement: Tuple[float, float]) -> Tuple[float, float]:
        """
        使用观测值更新状态估计
        
        结合预测值和观测值，得到更准确的状态估计。
        
        参数：
            measurement: 观测到的位置 (x, y)
            
        返回：
            更新后的位置 (x, y)
        """
        if not self.is_initialized:
            raise ValueError("滤波器未初始化，请先调用initialize方法")
            
        # 观测向量
        z = np.array([measurement[0], measurement[1]], dtype=np.float32)
        
        # 计算Kalman增益
        S = self.H @ self.covariance @ self.H.T + self.R  # 创新协方差
        K = self.covariance @ self.H.T @ np.linalg.inv(S)  # Kalman增益
        
        # 计算创新（残差）
        innovation = z - self.H @ self.state
        
        # 状态更新
        self.state = self.state + K @ innovation
        
        # 协方差更新
        I = np.eye(self.state_dim)
        self.covariance = (I - K @ self.H) @ self.covariance
        
        return (self.state[0], self.state[1])
        
    def get_state(self) -> Tuple[float, float, float, float]:
        """
        获取当前状态
        
        返回：
            当前状态 (x, y, vx, vy)
        """
        return (self.state[0], self.state[1], self.state[2], self.state[3])
        
    def get_position(self) -> Tuple[float, float]:
        """
        获取当前位置
        
        返回：
            当前位置 (x, y)
        """
        return (self.state[0], self.state[1])
        
    def get_velocity(self) -> Tuple[float, float]:
        """
        获取当前速度
        
        返回：
            当前速度 (vx, vy)
        """
        return (self.state[2], self.state[3])
        
    def get_position_uncertainty(self) -> float:
        """
        获取位置估计的不确定性
        
        返回位置协方差矩阵的迹，用于评估估计的可信度。
        
        返回：
            位置不确定性（标准差）
        """
        position_covariance = self.covariance[:2, :2]
        return np.sqrt(np.trace(position_covariance))
        
    def set_process_noise(self, noise_std: float) -> None:
        """
        设置过程噪声标准差
        
        动态调整过程噪声，适应不同的运动模式。
        
        参数：
            noise_std: 过程噪声标准差
        """
        self.process_noise = noise_std
        
        # 重新计算过程噪声协方差矩阵
        self.Q = np.array([
            [self.dt**4/4, 0, self.dt**3/2, 0],
            [0, self.dt**4/4, 0, self.dt**3/2],
            [self.dt**3/2, 0, self.dt**2, 0],
            [0, self.dt**3/2, 0, self.dt**2]
        ], dtype=np.float32) * self.process_noise**2
        
    def set_measurement_noise(self, noise_std: float) -> None:
        """
        设置测量噪声标准差
        
        根据观测精度调整测量噪声。
        
        参数：
            noise_std: 测量噪声标准差
        """
        self.measurement_noise = noise_std
        self.R = np.eye(self.measurement_dim, dtype=np.float32) * self.measurement_noise**2
        
    def reset(self) -> None:
        """
        重置滤波器状态
        
        清除所有状态信息，准备重新初始化。
        """
        self.state = np.zeros(self.state_dim)
        self.covariance = np.eye(self.state_dim) * 1000
        self.is_initialized = False 